package com.study.spark.scala.dataset

import org.apache.spark.sql.SparkSession

/**
 * 列操作
 *
 * @author stephen
 * @date 2019-09-27 16:42
 */
object ColumnDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .appName("ColumnDemo")
      .master("local[*]")
      .getOrCreate()

    // 隐式转换
    import org.apache.spark.sql.functions._
    import spark.implicits._

    val df = spark.createDataset(
      Seq((3, "zhangsan", "CN", 20),
        (4, "lisi", null, 30),
        (5, null, "JP", 40)))
      .toDF("id", "name", "country", "age")

    // 新增列，使用lit()设置初始化值
    df.withColumn("new_col", lit("new_str_col")).show()
    df.withColumn("new_col", $"age" + 1).show()
    df.withColumn("new_col", expr("age+1")).show()

    // 删除列
    df.drop($"age").show()
    df.drop("age", "country").show()

    // 得到所有列名 df.columns
    println(df.columns.toList)

    // null值操作
    // 把country列中所有的null值替换成'CN'
    df.na.fill("CN", Seq("country")).show()
    // 一次修改不同类型的多个列
    df.na.fill(Map(("country" -> "CN"), ("name" -> "wangwu"))).show()
    // 过滤掉含有null值的行
    df.na.drop(Seq("name", "country")).show()
    // 过滤掉指定的若干列中，有效值少于n列的行
    df.na.drop(2, Seq("name", "country")).show()
    // 如果不指定列名列表，则默认列名列表就是所有列
    df.na.drop(2).show()

    spark.stop()
  }
}
